# -*- coding: utf-8 -*-
"""
Created on Wed Apr 17 20:33:16 2013
@author: Maxim
"""
import numpy as np
import aircraft
import aerodynamics
import matplotlib.pyplot as plt
import matplotlib
import FlightConditions as FC
import mass
from warnings import warn
import report_tools as rt

class StabilityReport(aerodynamics.aero3d_VLM):
    """
    Generates plots of stability coefficients vs. Velocity or Mach number. 
    All required coefficients can be specified explicitly that makes code 
    very flexible. Center of gravity and mass can be specified explicitly as 
    well as through mass list.

    Examples
    --------

    importing necessary modules
    
    >>> import aircraft
    >>> import mass
    >>> from stability_report import StabilityReport
    
    load aircraft configuration
    
    >>> ac = aircraft.load('V05')
    
    create mass list with payload
    
    >>> newMassList = mass.MassList('sample mass')
    
    create required mass components
    
    >>> panel     = mass.MassComponent('panel', 10.0,[1.4,0.0,0.1])
    >>> misc      = mass.MassComponent('misc',  52.0,[2.32,0.0,0.0])
    >>> passenger = mass.MassComponent('passenger', 85.0,[1.86,0.25,-0.075])
    >>> baggage   = mass.MassComponent('baggage', 75.0,[2.85,0,-0.125])
    
    add mass components to mass list

    >>> newMassList.add_component(panel)
    >>> newMassList.add_component(misc)
    >>> newMassList.add_component(passenger)
    >>> newMassList.add_component(baggage)

    initialize new stability report
    
    >>> sr = StabilityReport(ac)

    add flight condition
    
    >>> sr.add_flight_condition(1500.0,[0.08,0.2,5],True)
    
    add required parameters

    >>> sr.set_required_param('coefficients',['CL','CD'], ['lift coefficient','drag coefficient'])
    >>> sr.set_required_param('dutch roll',['w','Z'], ['dutch roll nat.freq.','dutch roll damp.ratio'])

    generate stability plots with given parameters

    >>> sr.generate()
    """
    def init_report(self):
        """
        Reinitializing report generation.
        """
        self.trim = True
        self.alphaSweep = False
        self.reportStaticMargin = False
        self.reprtHingeMoments = False #otherwise hinge coefficients are displayed
        self.parBasic   = list()
        self.parCoef    = list()
        self.parDeriv   = list()
        self.parPhugoid = list()
        self.parSPeriod = list()
        self.parDutchRoll=list()
        self.parSpiral  = list()
        self.parRoll    = list()
        self.parHinge   = list()
        self.labelBasic   = list()
        self.labelCoef    = list()
        self.labelDeriv   = list()
        self.labelPhugoid = list()
        self.labelSPeriod = list()
        self.labelDutchRoll=list()
        self.labelSpiral  = list()
        self.labelRoll    = list()
        self.labelHinge   = list()
        self.rsltBasic   = list()
        self.rsltCoef    = list()
        self.rsltDeriv   = list()
        self.rsltPhugoid = list()
        self.rsltSPeriod = list()
        self.rsltDutchRoll=list()
        self.rsltSpiral  = list()
        self.rsltRoll    = list()
        self.rsltHinge   = list()
        self.dynamicPressure = list()
        self.cgList      = list()
        self.massList    = list()
        self.loadFactorList = list()
        self.legendCG    = list()
        self.useMach = False
        self.plot = True
        self.report = True
    def add_flight_condition(self,altitude,machSeq,
                             CmTrim=0.0,aileron=0.0,rudder=0.0,flap=0.0):
        """
        add flight conditions in stack
        
        Parameters
        ----------
        
        altitude : float, m
            altitude at which analysis will be performed
        machSeq : 3d list, m/sec
            velocity sequence in format [start velocity, end velocity, number of points]
        useMach : bool
            if True then machSeq is treated as mach number sequence otherwise 
            treated as velocity in m/sec sequence
        loadFactor : float
            load factor
        CmTrim : float
            trim pitch moment coefficient
        """
        self.altitude = altitude
        atm = FC.ISAtmosphere_simple(altitude)
        if type(machSeq)==type(float()):
            machSeq = np.array([machSeq])
        else:
            machSeq = np.linspace(machSeq[0],machSeq[1],machSeq[2])
        if machSeq[0]<1.0:
            self.useMach=True
        else:
            self.useMach=False
        if self.useMach:
            self.Mach = machSeq
            self.velocity = machSeq * atm.soundSpeed
            self.xlabel = 'Mach number'
        else:
            self.velocity = machSeq
            self.Mach = machSeq / atm.soundSpeed
            self.xlabel = 'Velocity, m/sec'
        self.trim = True
        self.density = atm.density
        self.CmTrim = CmTrim
        self.aileron = aileron
        self.rudder = rudder
        self.flap = flap
    def set_alpha_sweep(self,altitude,velocity,alpha=None,beta=0.0,
                        elevator=0.0,aileron=0.0,rudder=0.0,flap=0.0):
        """
        set alpha sweep analysis at fixed beta. 
        Alpha is limited to linear region only
        """
        self.trim = False
        self.alphaSweep = True
        self.xlabel = 'alpha, deg'
        if alpha==None:
            alpha = np.array([-2.0,0.0,2.0,4.0,6.0])
        self.alpha = alpha
        self.beta = beta
        self.velocity = velocity
        self.altitude = altitude
        atm = FC.ISAtmosphere_simple(altitude)
        self.density = atm.density
        self.elevator = elevator
        self.aileron  = aileron
        self.rudder   = rudder
        self.flap     = flap
    def set_beta_sweep(self,altitude,velocity,alpha=0.0,beta=None,
                       elevator=0.0,aileron=0.0,rudder=0.0,flap=0.0):
        """
        set beta sweep analysis at fixed alpha.
        """
        self.trim = False
        self.alphaSweep = False
        self.xlabel = 'beta, deg'
        if beta==None:
            beta = np.array([0.0,2.0,4.0,6.0,8.0])
        self.alpha = alpha
        self.beta = beta
        self.velocity = velocity
        self.altitude = altitude
        atm = FC.ISAtmosphere_simple(altitude)
        self.density = atm.density
        self.elevator = elevator
        self.aileron  = aileron
        self.rudder   = rudder
        self.flap     = flap
    def set_required_param(self,category,paramList,paramLabels=[]):
        """
        Set parameters required to plot. All parameters will be plotted vs. 
        velocity or Mach number for different load cases.
        
        Parameters
        ----------
        
        category : string
            name of the category at which coefficient is defined:
                - basic
                - coefficients
                - derivatives
                - phugoid
                - short period
                - dutch roll
                - roll
                - hinge moment
        paramList : string list
            list of parameters in given category. (for example "alpha", "xNP" etc.)
        paramLabels : string list
            list of labels which will be displayed as ylabel.
        """
        if len(paramList)!=len(paramLabels):
            paramLabels=paramList
        if category=='basic':
            self.parBasic += paramList
            self.labelBasic += paramLabels
        elif category=='coefficients':
            self.parCoef += paramList
            self.labelCoef += paramLabels
        elif category=='derivatives':
            self.parDeriv += paramList
            self.labelDeriv += paramLabels
        elif category=='phugoid':
            self.parPhugoid += paramList
            self.labelPhugoid += paramLabels
        elif category=='short period':
            self.parSPeriod += paramList
            self.labelSPeriod += paramLabels
        elif category=='dutch roll':
            self.parDutchRoll += paramList
            self.labelDutchRoll += paramLabels
        elif category=='short period':
            self.parSpiral += paramList
            self.labelSpiral += paramLabels
        elif category=='roll':
            self.parRoll += paramList
            self.labelRoll += paramLabels
        elif category=='hinge moment':
            self.parHinge += paramList
            self.labelHinge += paramLabels
        else:
            Warning('required parameter category is not found')
    def add_payload_mass_list(self,massList):
        """
        Replaces aircraft.mass.payload MassList with massList so total mass and 
        CG are recalculated for given massList
        
        Parameters
        ----------
        
        massList : MassList
            mass.MassList object that contains all payload information 
            for example: pilot, passenger, baggage etc.
        """
        self.legendCG.append(massList.name)
        self.aircraft.mass.payload = massList
        newCG   = self.aircraft.get_CG(True)
        newMass = self.aircraft.get_mass_total()
        self.loadFactorList.append(self.loadFactor)
        self.cgList.append(newCG)
        self.massList.append(newMass)
    def add_cg(self,CG,mass,loadFactor=1.0,legend=''):
        """
        Adds aircraft CG and mass explicitly.
        
        Parameters
        ----------
        
        CG : 3d float array, mm
            aircraft center of gravity
        mass : float, kg
            aircraft mass
        loadFactor : float
            load factor
        legend : string
            name of the load case
        """
        self.cgList.append(CG)
        self.massList.append(mass)
        self.loadFactorList.append(loadFactor)
        self.legendCG.append(legend)
    def generate(self,fontSize=None):
        """
        generates stability plots with parameters defined.
        """
        self._analyze()
        self._extract_data()
        if self.reportStaticMargin:
            self._calc_sm()
        if self.report:
            self._report_results()
        if self.plot:
            self._plot_results(fontSize)
    def _calc_sm(self):
        c = self.aircraft.wing.MAC
        SM = np.zeros([self.nCG,self.nMach])
        for i,cg in enumerate(self.cgList):
            for j in range(self.nMach):
                k = i*self.nMach + j
                NP = self.rsltBasic[k].__dict__['xNP']
                SM[i,j] = (NP-cg[0]) / c
        self.dataToPlot.append(SM)
        self.labelToPlot.append('Static margin')
    def _report_results(self):
        rpt = rt.Report()
        for i,case in enumerate(self.cgList):
            rpt.new_line()
            rpt.add_mass_line1(self.legendCG[i],self.massList[i],case,True)
            if self.useMach:
                rpt.add_stability_line(self.xlabel,self.Mach,2)
            else:
                rpt.add_stability_line(self.xlabel,self.velocity,2)
            for j in range(len(self.dataToPlot)):
                data = self.dataToPlot[j][i]
                label = self.labelToPlot[j]
                rpt.add_stability_line(label,data)
        print rpt.out

    def _analyze(self):
        def append_results(rslt):
            self.rsltBasic.append(rslt)
            self.rsltCoef.append(rslt.coefficients)
            self.rsltDeriv.append(rslt.derivs)
            self.rsltPhugoid.append(rslt.dynamicStability.phugoid)
            self.rsltSPeriod.append(rslt.dynamicStability.shortPeriod)
            self.rsltDutchRoll.append(rslt.dynamicStability.dutchRoll)
            self.rsltSpiral.append(rslt.dynamicStability.spiral)
            self.rsltRoll.append(rslt.dynamicStability.roll)
            self.rsltHinge.append(rslt.hingeMoments)
        self.nCG = len(self.cgList)
        I = self.aircraft.get_inertia()
        for mass1,cg,n in zip(self.massList,self.cgList,self.loadFactorList):
            if self.trim:
                self.nMach = len(self.velocity)
                for j,V in enumerate(self.velocity):
                    CD0 = self.aircraft.get_drag(V,self.altitude)
                    fc = aerodynamics.flightConditions()
                    fc.addTrimmedFlightCondition('stabReport',mass1,cg,I,V,self.density,
                                                 n,CD0,self.CmTrim,self.flap,
                                                 self.aileron,self.rudder)
                    append_results(self.runVLM(fc).results[0])
            else:
                CD0 = self.aircraft.get_drag(self.velocity,self.altitude)
                if self.alphaSweep:
                    self.nMach = len(self.alpha)
                    for alpha in self.alpha:
                        fc = aerodynamics.flightConditions()
                        fc.add_flight_condition('untrimmed',mass1,cg,I,self.velocity,self.density,n,
                                                CD0,alpha,self.beta,self.flap,self.elevator,self.aileron,self.rudder)
                        append_results(self.runVLM(fc).results[0])
                else:
                    self.nMach = len(self.beta)
                    for beta in self.beta:
                        fc = aerodynamics.flightConditions()
                        fc.add_flight_condition('untrimmed',mass1,cg,I,self.velocity,self.density,n,
                                                CD0,self.alpha,beta,self.flap,self.elevator,self.aileron,self.rudder)
                        append_results(self.runVLM(fc).results[0])
    def _extract_data_single(self,rslt,paramList,labelList):
        data = list()
        label = list()
        for ii,param in enumerate(paramList):
            data1 = np.zeros([self.nCG,self.nMach])
            for i in range(self.nCG):
                for j in range(self.nMach):
                    k = i*self.nMach + j
                    data1[i,j] = rslt[k].__dict__[param]
            data.append(data1)
            label.append(labelList[ii])
        return data,label
    def get_data(self,category,name):
        if category=='basic':
            rslt = self.rsltBasic
        elif category=='coefficients':
            rslt = self.rsltCoef
        elif category=='derivatives':
            rslt = self.rsltDeriv
        elif category=='phugoid':
            rslt = self.rsltPhugoid
        elif category=='short period':
            rslt = self.rsltSPeriod
        elif category=='dutch roll':
            rslt = self.rsltDutchRoll
        elif category=='spiral':
            rslt = self.rsltSpiral
        elif category=='roll':
            rslt = self.rsltRoll
        elif category=='hinge moment':
            rslt = self.rsltHinge
        data = np.zeros([self.nCG,self.nMach])
        for i in range(self.nCG):
            for j in range(self.nMach):
                k = i*self.nMach + j
                data[i,j] = rslt[k].__dict__[name]
        return data
    def _extract_data(self):
        self.dataToPlot = list()
        self.labelToPlot = list()
        if self.parBasic!=[]:
            extractedData = self._extract_data_single(self.rsltBasic,self.parBasic,self.labelBasic)
            self.dataToPlot += extractedData[0]
            self.labelToPlot += extractedData[1]
        if self.parCoef!=[]:
            extractedData = self._extract_data_single(self.rsltCoef,self.parCoef,self.labelCoef)
            self.dataToPlot += extractedData[0]
            self.labelToPlot += extractedData[1]
        if self.parDeriv!=[]:
            extractedData = self._extract_data_single(self.rsltDeriv,self.parDeriv,self.labelDeriv)
            self.dataToPlot += extractedData[0]
            self.labelToPlot += extractedData[1]
        if self.parPhugoid!=[]:
            extractedData = self._extract_data_single(self.rsltPhugoid,self.parPhugoid,self.labelPhugoid)
            self.dataToPlot += extractedData[0]
            self.labelToPlot += extractedData[1]
        if self.parSPeriod!=[]:
            extractedData = self._extract_data_single(self.rsltSPeriod,self.parSPeriod,self.labelSPeriod)
            self.dataToPlot += extractedData[0]
            self.labelToPlot += extractedData[1]
        if self.parDutchRoll!=[]:
            extractedData = self._extract_data_single(self.rsltDutchRoll,self.parDutchRoll,self.labelDutchRoll)
            self.dataToPlot += extractedData[0]
            self.labelToPlot += extractedData[1]
        if self.parSpiral!=[]:
            extractedData = self._extract_data_single(self.rsltSpiral,self.parSpiral,self.labelSpiral)
            self.dataToPlot += extractedData[0]
            self.labelToPlot += extractedData[1]
        if self.parRoll!=[]:
            extractedData = self._extract_data_single(self.rsltRoll,self.parRoll,self.labelRoll)
            self.dataToPlot += extractedData[0]
            self.labelToPlot += extractedData[1]
        if self.parHinge!=[]:
            extractedData = self._extract_data_single(self.rsltHinge,self.parHinge,self.labelHinge)
            self.dataToPlot += extractedData[0]
            self.labelToPlot += extractedData[1]

    def _plot_results(self,fontsize=None):
        if not fontsize==None:
            font = {'family':'sans-serif','weight':'normal','size':fontsize}
            matplotlib.rc('font',**font)
        for i,data in enumerate(self.dataToPlot):
            yParamLabel = self.labelToPlot[i]
            plt.figure(i+1)
            plt.hold(True)
            plt.grid(True)
            plt.xlabel(self.xlabel)
            plt.ylabel(yParamLabel)
            for j in range(self.nCG):
                if self.trim:
                    if self.useMach:
                        plt.plot(self.Mach,data[j,:])
                    else:
                        plt.plot(self.velocity,data[j,:])
                else:
                    if self.alphaSweep:
                        plt.plot(self.alpha,data[j,:])
                    else:
                        plt.plot(self.beta,data[j,:])
            plt.legend(self.legendCG,0)
        plt.show()

def generate_payload_mass_list():
    CGlist = list()

    passenger  = mass.MassComponent('passenger',85.0,[1.86,0.25,-0.075])
    pilotHeavy = mass.MassComponent('pilot',    85.0,[1.86,-0.25,-0.075])
    pilotLight = mass.MassComponent('pilot',    55.0,[1.86,-0.25,-0.075])
    baggage    = mass.MassComponent('baggage',  75.0,[2.85,0,-0.125])
    panel      = mass.MassComponent('panel',    10.0,[1.4,0.0,0.1])
    misc       = mass.MassComponent('misc',     52.0,[2.32,0.0,0.0])

    CGfw = mass.MassList('forward CG')
    CGfw.add_component(pilotHeavy)
    CGfw.add_component(passenger)
    CGfw.add_component(panel)
    CGfw.add_component(misc)
    
    CGaft = mass.MassList('aft CG')
    CGaft.add_component(pilotLight)
    CGaft.add_component(panel)
    CGaft.add_component(misc)
    CGaft.add_component(baggage)
    
    CGmid = mass.MassList('middle CG')
    CGmid.add_component(pilotHeavy)
    CGmid.add_component(panel)
    CGmid.add_component(misc)

    CGlist.append(CGfw)
    CGlist.append(CGmid)
    CGlist.append(CGaft)
    return CGlist




def test_report():
    ac = aircraft.load('V0510')
    print ac.wing.MAC
    print ac.vStab.MAC
    print ac.hStab.MAC
    sr = StabilityReport(ac)
    sr.init_report()
    CG = generate_payload_mass_list()
    sr.add_flight_condition(0.0,[20.0,30.0,10],False,flap=15.0)
    for cg in CG:
        sr.add_payload_mass_list(cg)
        print cg.display()
#    sr.reportStaticMargin = True
#    sr.reprtHingeMoments = True
#    sr.set_required_param('basic',['alpha','elevator'],['Trim AOA, deg','elev. deflection,deg'])
#    sr.set_required_param('coefficients',['CL','CD'],['lift coefficient','drag coefficient'])
#    sr.set_required_param('derivatives',['Cma','Cnb','Clb'],['Cm-alpha','Cn-beta','Cl-beta'])

    sr.set_required_param('short period',['w','Z'],['short period nat.freq.','short period damp.ratio'])
    sr.set_required_param('dutch roll',['w','Z'],['dutch roll nat.freq.','dutch roll damp.ratio'])
#    sr.set_required_param('hinge moment',['elevator','flap'],['elev hinge moment coeficient','flap hinge moment coefficient'])
    sr.generate()

def new_report():
    ac = aircraft.load('V0510')
    sr = StabilityReport(ac)
    sr.init_report()
    sr.plot = True
    sr.report = False
    sr.add_cg([1.711,-0.040,-0.124],593.1,1.0,'aft load')
    sr.add_cg([1.824,-0.0,-0.103],583.1,1.0,'fwd load')
    sr.add_flight_condition(1500.0,[30.0,70.0,7])
    #sr.set_alpha_sweep(0.0,70.0,aileron=20.0)
    sr.reportStaticMargin = True
#    sr.set_required_param('basic',['xNP','alpha','elevator'],['neutral point','trim AOA','trim elevator defl.'])
#    sr.set_required_param('coefficients',['CL','CD'],['lift coef.','drag coef.'])
#    sr.set_required_param('derivatives',['Clb','Cnb','Cma'],['Cl beta','Cn beta','Cm alpha'])
#    sr.set_required_param('short period',['Wn','Z'],['Short period nat freq.','Short period damping ratio'])
#    sr.set_required_param('dutch roll',['Wn','Z'],['Dutch roll nat freq.','Dutch roll damping ratio'])
    sr.set_required_param('hinge moment',['Mrudder','Maileron','aileron','Melevator'],['rudder, N*m','aileron, N*m','aileron coef','elevator, N*m'])
    sr.generate(14)

def v200_stability():
    # report for KLA100 CDR 2014-03-31
    ac = aircraft.load('V204-ext_aileron')
    #ac.display()
    sr = StabilityReport(ac)
    sr.init_report()
    sr.plot = True
    wingLoc = ac.wing.aapex[0]
    wingMAC = ac.wing.MAC
    sr.add_cg([0.2513*wingMAC+wingLoc, 0,0],620,1.0,'maximum mass,fw CG')
    sr.add_cg([0.3691*wingMAC+wingLoc, 0,0],620.,1.0,'maximum mass, aft CG')
    sr.add_cg([0.1649*wingMAC+wingLoc, 0,0],403.5,1.0,'minimum mass, abs fw CG')
    sr.add_flight_condition(1500.0,[30,80,10])
    sr.reportStaticMargin = True
    sr.set_required_param('basic',['xNP','alpha','elevator'],['neutral point','trim AOA','trim elevator defl.'])
    sr.set_required_param('coefficients',['CL','CD'],['lift coef.','drag coef.'])
    sr.set_required_param('derivatives',['Clb','Cnb','Cma'],['Cl beta','Cn beta','Cm alpha'])
    sr.set_required_param('short period',['Wn','Z'],['Short period nat freq.','Short period damping ratio'])
    sr.set_required_param('dutch roll',['Wn','Z'],['Dutch roll nat freq.','Dutch roll damping ratio'])
    sr.generate()

if __name__=="__main__":
    v200_stability()